import { VmoRequest } from 'vmo-request';
import { MilvusClient, DataType } from '@zilliz/milvus2-sdk-node';

// 定义请求和响应的类型
interface EmbeddingRequest {
  model: string;
  prompt: string;
}

interface EmbeddingResponse {
  data: {
    embedding: number[];
    model: string;
    created_at: string;
  }
}

// 创建 Milvus 客户端
const milvusClient = new MilvusClient({
  address: 'localhost:19530',
});
// 创建请求实例
const reRankingModel = new VmoRequest({
  baseURL: 'https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank',
  timeout: 5000, // 30秒超时，因为模型推理可能需要较长时间
  headers:{
    "Authorization": "Bearer sk-1e36e14fe6fa45c5a8f314c408f1ae39",
    "Content-Type":"application/json",
  }
}, true); // 创建伴随的 axios 实例

// 创建请求实例
const vmoRequest = new VmoRequest({
  baseURL: 'http://localhost:11434',
  timeout: 5000, // 30秒超时，因为模型推理可能需要较长时间
}, true); // 创建伴随的 axios 实例

async function getReranking(query:string,documents:string[]){
  console.log(query,documents)
  const reranking = await reRankingModel.request({
    method:'post',
    data:{
      model:'gte-rerank-v2',
      input:{
        query,
        documents
      },
      "parameters": {
          "return_documents": true,
          "top_n": 5
      }
    }
  })
  return reranking
}
/**
 * 初始化 Milvus 集合
 */
async function initMilvusCollection() {
  try {
    // 检查集合是否存在
    const hasCollection = await milvusClient.hasCollection({
      collection_name: 'text_embeddings',
    });

    if (!hasCollection.value) {      
      // 创建集合
      await milvusClient.createCollection({
        collection_name: 'text_embeddings',
        fields: [
          {
            name: 'id',
            data_type: DataType.Int64,
            is_primary_key: true,
            autoID: true,
          },
          {
            name: 'text',
            data_type: DataType.VarChar,
            max_length: 65535,
          },
          {
            name: 'embedding',
            data_type: DataType.FloatVector,
            dim: 1024, // bge-m3 的向量维度
          },
        ],
      });

      // 创建索引
      await milvusClient.createIndex({
        collection_name: 'text_embeddings',
        field_name: 'embedding',
        index_type: 'HNSW',
        metric_type: 'COSINE',
        params: { M: 8, efConstruction: 64 },
      });

      console.log('Milvus 集合创建成功');
    } else {
      // milvusClient.dropCollection({collection_name:'text_embeddings'})
      console.log('Milvus 集合已存在');
    }

    // 加载集合到内存
    await milvusClient.loadCollection({
      collection_name: 'text_embeddings',
    });
    console.log('Milvus 集合加载成功');
  } catch (error) {
    console.error('初始化 Milvus 集合失败:', error);
    throw error;
  }
}

/**
 * 使用 bge-m3 模型生成文本嵌入
 * @param text 需要生成嵌入的文本
 * @returns 文本的嵌入向量
 */
async function generateEmbedding(text: string): Promise<number[]> {
  try {
    const response = await vmoRequest.request({
      url: '/api/embeddings',
      method: 'post',
      data: {
        model: 'bge-m3',
        prompt: text,
      } as EmbeddingRequest
    }, 1, 3); // 不缓存，失败重试3次
    // console.log(response)
    // 打印响应数据（避免循环引用）
    console.log('API 响应数据:', {
      embedding: response?.data?.embedding.length,
    });

    // 验证响应数据
    if (!response || typeof response !== 'object') {
      throw new Error('响应数据格式错误：不是有效的对象');
    }

    const embeddingResponse = response as EmbeddingResponse;
    if (!embeddingResponse?.data?.embedding || !Array.isArray(embeddingResponse?.data?.embedding)) {
      throw new Error('响应数据格式错误：embedding 字段不是有效的数组');
    }

    return embeddingResponse?.data?.embedding;
  } catch (error) {
    console.error('生成嵌入失败:', error);
    throw new Error(`生成文本嵌入失败: ${error instanceof Error ? error.message : '未知错误'}`);
  }
}

/**
 * 存储文本和其嵌入向量到 Milvus
 * @param text 文本内容
 * @param embedding 文本的嵌入向量
 */
async function storeEmbedding(text: string, embedding: number[]) {
  try {
    await milvusClient.insert({
      collection_name: 'text_embeddings',
      fields_data: [
        {
          text,
          embedding,
        },
      ],
    });
    console.log('向量数据存储成功');
  } catch (error) {
    console.error('存储向量数据失败:', error);
    throw error;
  }
}

/**
 * 向量检索
 * @param queryText 查询文本
 * @param topK 返回最相似的前K个结果
 * @returns 检索结果
 */
async function searchSimilarVectors(queryText: string, topK: number = 1) {
  try {
    // 确保集合已加载
    await milvusClient.loadCollection({
      collection_name: 'text_embeddings',
    });

    // 生成查询文本的向量
    const queryEmbedding = await generateEmbedding(queryText);
    
    // 执行向量检索
    const searchResults = await milvusClient.search({
      collection_name: 'text_embeddings',
      vector: queryEmbedding,
      output_fields: ['text','id'],
      limit: topK,
      metric_type: 'COSINE',
    });

    return searchResults;
  } catch (error) {
    console.error('向量检索失败:', error);
    throw error;
  }
}

// 使用示例
export async function ragSample() {
  try {
    // 初始化 Milvus 集合
    await initMilvusCollection();
    const query = "母亲"
    // const text = `晨光斜进厨房纱窗时，空气里已浮动着煎蛋和米粥暖甜的气息。梳齿穿过我的发梢时轻柔得像风梳过柳条。油锅里的滋滋声，暖水壶吐出的白汽，还有她俯身为我系上红领巾时发梢掠过的淡淡馨香，织成一张柔韧的网。桌上摆着切好的苹果块，每一粒都削圆了棱角。放学路上远远望见那个熟悉的影子守在校门，雨伞下意识地向我倾斜，任雨水顺着伞骨打湿一侧肩头，脸上却绽开盈盈笑意。她的手粗糙却温厚，有时拂过我额角，像一片承载着阳光的叶子轻轻停驻。光阴静流如窗台上那盆茉莉的暗香，虽无言，却早已悄然缝进每一寸衣襟、系紧每一粒纽扣、刻入血脉纹路的缝隙里。`;
    // const embedding = await generateEmbedding(text);
    
    // 验证嵌入向量
    // if (!embedding || !Array.isArray(embedding)) {
    //   throw new Error('生成的嵌入向量无效');
    // }

    // console.log('生成的嵌入向量:', embedding);
    // console.log('向量维度:', embedding.length);

    // 存储向量数据
    // await storeEmbedding(text, embedding);

    // 测试向量检索
    console.log('\n开始测试向量检索...');
    const searchResults = await searchSimilarVectors(query,10);
    const rerankingData = await getReranking(query,searchResults?.results.map(item=>item.text))
    console.log('检索结果:', JSON.stringify(searchResults, null, 2));
    console.log('重排结果:',query ,rerankingData.data.output.results.map((item:Record<string,any>)=>{
      item.document = item.document.text;
      return item;
    }),rerankingData.data);
  } catch (error) {
    console.error('程序执行失败:', error instanceof Error ? error.message : '未知错误');
  }
}
